﻿using System.Collections.Generic;
using UnityEngine;
using Unity.Entities;
using Unity.Mathematics;
using Unity.Transforms;
using Unity.Collections;
using Unity.Jobs;
using Unity.Burst;
using Unity.Rendering;
using Unity.Profiling;
using Jinndev.ECS;

namespace Jinndev.Rendering2D {

    /// <summary>
    /// 需要组件：CustomRenderer, LocalToWorld, AtlasTilingOffset
    /// </summary>
    [UpdateInGroup(typeof(PresentationSystemGroup))]
    //[UpdateInGroup(typeof(SimulationSystemGroup), OrderLast = true)]
    //[AlwaysSynchronizeSystem]
    public class CustomRenderSystem : JobComponentSystem {

        private MaterialPropertyBlock properties = new MaterialPropertyBlock();
        private Matrix4x4[] matrices = new Matrix4x4[1023];
        private Vector4[] tilingOffsets = new Vector4[1023];
        private int drawCall = 0;

        private EntityQuery entityQuery;

        private static readonly ProfilerMarker profileOther = new ProfilerMarker("OtherSystemJob");
        private static readonly ProfilerMarker profileQuery = new ProfilerMarker("Query");
        private static readonly ProfilerMarker profileSort = new ProfilerMarker("Sort");
        private static readonly ProfilerMarker profileCombine = new ProfilerMarker("Combine");
        private static readonly ProfilerMarker profileAssembly = new ProfilerMarker("Assembly");
        private static readonly ProfilerMarker profileBatch = new ProfilerMarker("Batch");

        protected override void OnCreate() {
            entityQuery = GetEntityQuery(
                ComponentType.ReadOnly<CustomRenderer>(),
                ComponentType.ReadOnly<LocalToWorld>(),
                ComponentType.ReadOnly<AtlasTilingOffset>(),
                ComponentType.Exclude<Culled>()
            );
        }

        [BurstCompile]
        struct AssemblyJob : IJobParallelFor {
            [WriteOnly] public NativeArray<Matrix4x4> matrixOutput;
            [WriteOnly] public NativeArray<Vector4> atlasTilingOffsetOutput;
            [ReadOnly] public NativeArray<AtlasTilingOffset> atlasTilingOffsetInput;
            [ReadOnly] public NativeArray<LocalToWorld> localToWorldInput;

            public void Execute(int index) {
                matrixOutput[index] = localToWorldInput[index].Value;
                atlasTilingOffsetOutput[index] = atlasTilingOffsetInput[index].value;
            }
        }

        // Chunk比较器
        struct ChunkComparer : IComparer<ArchetypeChunk> {
            public SharedComponentTypeHandle<CustomRenderer> typeCustomRenderer;

            public int Compare(ArchetypeChunk x, ArchetypeChunk y) {
                return x.GetSharedComponentIndex(typeCustomRenderer).CompareTo(y.GetSharedComponentIndex(typeCustomRenderer));
            }
        }

        protected override JobHandle OnUpdate(JobHandle inputDeps) {
            drawCall = 0;

            profileOther.Begin();
            inputDeps.Complete();
            profileOther.End();

            profileQuery.Begin();
            var typeCustomRenderer = GetSharedComponentTypeHandle<CustomRenderer>();
            var typeAtlasTilingOffset = GetComponentTypeHandle<AtlasTilingOffset>();
            var typeLocalToWorld = GetComponentTypeHandle<LocalToWorld>();

            NativeArray<ArchetypeChunk> chunkArray = entityQuery.CreateArchetypeChunkArrayAsync(Allocator.TempJob, out inputDeps);
            inputDeps.Complete();

            int entityCount = entityQuery.CalculateEntityCount();
            profileQuery.End();


            // chunk排序，相同Renderer的放一起
            profileSort.Begin();
            chunkArray.Sort(new ChunkComparer() { typeCustomRenderer = typeCustomRenderer });
            profileSort.End();


            // 所有chunk种的Component取出来放一起
            profileCombine.Begin();
            NativeArray<LocalToWorld> localToWorldArray = new NativeArray<LocalToWorld>(entityCount, Allocator.TempJob);
            NativeArray<AtlasTilingOffset> atlasTilingOffsetArray = new NativeArray<AtlasTilingOffset>(entityCount, Allocator.TempJob);
            int cacheCount = 0;

            for (int i = 0; i < chunkArray.Length; i++) {
                ArchetypeChunk chunk = chunkArray[i];
                NativeArray<AtlasTilingOffset> tilingOffsetCache = chunk.GetNativeArray(typeAtlasTilingOffset);
                NativeArray<LocalToWorld> localToWorldCache = chunk.GetNativeArray(typeLocalToWorld);

                NativeArray<AtlasTilingOffset>.Copy(tilingOffsetCache, 0, atlasTilingOffsetArray, cacheCount, chunk.Count);
                NativeArray<LocalToWorld>.Copy(localToWorldCache, 0, localToWorldArray, cacheCount, chunk.Count);
                cacheCount += chunk.Count;
            }
            profileCombine.End();


            // 取得value数据
            profileAssembly.Begin();
            NativeArray<Matrix4x4> matrixOutput = new NativeArray<Matrix4x4>(entityCount, Allocator.TempJob);
            NativeArray<Vector4> tilingOffsetOutput = new NativeArray<Vector4>(entityCount, Allocator.TempJob);

            AssemblyJob assemblyJob = new AssemblyJob() {
                matrixOutput = matrixOutput,
                atlasTilingOffsetOutput = tilingOffsetOutput,
                localToWorldInput = localToWorldArray,
                atlasTilingOffsetInput = atlasTilingOffsetArray
            };
            inputDeps = assemblyJob.Schedule(entityCount, 32, inputDeps);
            inputDeps.Complete();

            localToWorldArray.Dispose();
            atlasTilingOffsetArray.Dispose();
            profileAssembly.End();


            // 分批渲染
            profileBatch.Begin();
            int batchEntityCount = 0;
            int entityStartIndex = 0;
            CustomRenderer curCustomRenderer = default;

            for (int i = 0; i < chunkArray.Length; i++) {
                ArchetypeChunk chunk = chunkArray[i];
                CustomRenderer customRenderer = chunk.GetSharedComponentData(typeCustomRenderer, EntityManager);

                if (batchEntityCount > 0 && !customRenderer.Equals(curCustomRenderer)) {
                    while (batchEntityCount > 0) {
                        int batchCount = math.min(batchEntityCount, 1023);
                        NativeArray<Matrix4x4>.Copy(matrixOutput, entityStartIndex, matrices, 0, batchCount);
                        NativeArray<Vector4>.Copy(tilingOffsetOutput, entityStartIndex, tilingOffsets, 0, batchCount);
                        batchEntityCount -= batchCount;
                        entityStartIndex += batchCount;
                        Submit(batchCount, curCustomRenderer);
                    }
                }
                curCustomRenderer = customRenderer;
                batchEntityCount += chunk.Count;
            }

            if (batchEntityCount > 0) {
                while (batchEntityCount > 0) {
                    int batchCount = math.min(batchEntityCount, 1023);
                    NativeArray<Matrix4x4>.Copy(matrixOutput, entityStartIndex, matrices, 0, batchCount);
                    NativeArray<Vector4>.Copy(tilingOffsetOutput, entityStartIndex, tilingOffsets, 0, batchCount);
                    batchEntityCount -= batchCount;
                    entityStartIndex += batchCount;
                    Submit(batchCount, curCustomRenderer);
                }
            }
            profileBatch.End();

            matrixOutput.Dispose();
            tilingOffsetOutput.Dispose();
            chunkArray.Dispose();

            //Debug.Log(drawCall);
            return inputDeps;
        }

        private void Submit(int batchCount, CustomRenderer customRenderer) {
            if (customRenderer.mesh == null) {
                Debug.LogWarning("Mesh is null");
                return;
            }
            properties.SetVectorArray(RenderingConst.TilingOffsetShaderPropName, tilingOffsets);

            UnityEngine.Profiling.Profiler.BeginSample("Draw");
            try {
                Graphics.DrawMeshInstanced(customRenderer.mesh, 0, customRenderer.material, matrices, batchCount, properties);
                //Graphics.DrawMeshInstanced(customRenderer.mesh, 0, customRenderer.material, matrices, batchCount, properties, UnityEngine.Rendering.ShadowCastingMode.Off, false, 0, ViewManager.camera3D, UnityEngine.Rendering.LightProbeUsage.Off);
            }
            catch (System.Exception e) {
                Debug.LogError(e.Message + " \"" + customRenderer.material.name + "\"");
            }
            UnityEngine.Profiling.Profiler.EndSample();
            //properties.Clear();
            drawCall++;
        }

    }

}
